"""
This script holds the code for common data augmentation tools widely used
in deep learning for multispectral remote sensing imagery of shape [Height, Width, Bands].

@ Author: Bo Peng (bo.peng@wisc.edu)
@ Spatial Computing and Data Mining Lab, University of Wisconsin - Madison
@ Sponsor:
    - National Science Foundation (NSF)
    - Microsoft AI for Earth
@ Citation:
    - Peng, B., Huang, Q., & Rao, J. (2021, July). Spatiotemporal Contrastive Representation Learning
        for Building Damage Classification. In 2021 IEEE International Geoscience and Remote Sensing
        Symposium IGARSS (pp. 8562-8565). IEEE.
    - Peng, B., Huang, Q., Vongkusolkit, J., Gao, S., Wright, D. B., Fang, Z. N., & Qiang, Y. (2020).
        Urban flood mapping with bitemporal multispectral imagery via a self-supervised learning framework.
        IEEE Journal of Selected Topics in Applied Earth Observations and Remote Sensing, 14, 2001-2016.
© 2022 Bo Peng. All Rights Reserved.
"""

import math
import random
import cv2
import torch
from torchvision import transforms
import numpy as np


class RandomFlip(object):
    def __init__(self, p=0.5):
        """
        Random flipping the given image horizontally and vertically.

        Parameters
        ----------
        p: float
            Probability of random augmentation. Default: 0.5.
        """
        self.p = p

    def __call__(self, img):
        # horizontal
        if random.random() <= self.p:
            img = cv2.flip(img, 1)
        # vertical
        if random.random() <= self.p:
            img = cv2.flip(img, 0)
        return img

    def __repr__(self):
        return (
            f"Random flipping the image horizontally or "
            f"vertically with probability [p={self.p}]."
        )


class Resize(object):
    def __init__(self, size=(64, 64)):
        """
        Resize the given image to a specified size.

        Parameters
        ----------
        size: int or Tuple[int]
            Specified size of the resized image.
        """
        assert isinstance(size, (int, tuple))
        self.size = (size, size) if isinstance(size, int) else size

    def __call__(self, img):
        img = cv2.resize(img, self.size, interpolation=cv2.INTER_LINEAR)
        return img

    def __repr__(self):
        return "Resize image to {}.".format(self.size)


class RandomResizedCrop(object):
    def __init__(
        self, size=(64, 64), scale=(0.8, 1.0), ratio=(3.0 / 4, 4.0 / 3), num_trials=10
    ):
        """
        Randomly resize the given image and randomly crop a patch of a specified size.

        Parameters
        ----------
        size: int or Tuple[int]
            Specified size of the cropped image patch. Default: (64, 64).
        scale: Tuple[float]
            Scale range to which the given image will be randomly resized.
        ratio: Tuple[float]
            Ratio range to which the (height, width) of the given image will be randomly resized.
        num_trials: int
            Number of trials to attempt randomly resizing and cropping.
        """
        super(RandomResizedCrop, self).__init__()

        self.size = (size, size) if isinstance(size, int) else size
        self.scale = scale
        self.ratio = ratio
        self.num_trials = num_trials
        self.resize_sample = Resize(self.size)

    def __call__(self, img):

        for _ in range(self.num_trials):
            h, w = img.shape[:2]
            area = w * h
            target_area = random.uniform(self.scale[0], self.scale[1]) * area
            aspect_ratio = random.uniform(self.ratio[0], self.ratio[1])

            # aspect_ration = width / height, width * height = target_area
            h_crop = int(math.sqrt(target_area / aspect_ratio))
            w_crop = int(math.sqrt(target_area * aspect_ratio))

            # with probability of 0.5, aspect_ration = height / width
            if random.random() <= 0.5:
                w_crop, h_crop = h_crop, w_crop

            # if the cropped image size not exceed the original image
            if h_crop <= h and w_crop <= w:
                # sample the location of the left upper corner of the cropped image patch
                top = random.randint(0, h - h_crop)
                left = random.randint(0, w - w_crop)
                img = img[top : top + h_crop, left : left + w_crop]
                break

        img = cv2.resize(img, self.size, interpolation=cv2.INTER_LINEAR)

        return img

    def __repr__(self):
        return (
            f"Random Resized Crop: [size={self.size}] "
            f"[scale={self.scale}] [ratio={self.ratio}]."
        )


class RandomRotate(object):
    def __init__(self, angle=10):
        """Rotate the given image (around the image center) by a random degree,
        and crop the rotated image to the largest center patch without nodata at corners.

        Parameters
        ----------
        angle: float
            Range of degree (-d ~ +d) to which the given image will be rotated. Default: 10.
        """
        self.angle = angle

    def __call__(self, img):

        # sample rotation degree
        degree = np.random.uniform(-self.angle, self.angle)
        # ignore small rotations
        if np.abs(degree) <= 1.0:
            return img

        # get the max area rectangular within the rotated image
        # ref: stackoverflow.com/questions/16702966/rotate-image-and-crop-out-black-borders
        h, w = img.shape[:2]
        side_long = float(max([h, w]))
        side_short = float(min([h, w]))

        # since the solutions for angle, -angle and pi-angle are all the same,
        # it suffices to look at the first quadrant and the absolute values of sin,cos:
        sin_a = np.abs(np.sin(np.pi * degree / 180))
        cos_a = np.abs(np.cos(np.pi * degree / 180))

        if side_short <= 2.0 * sin_a * cos_a * side_long:
            # half constrained case: two crop corners touch the longer side,
            # the other two corners are on the mid-line parallel to the longer line
            x = 0.5 * side_short
            if w >= h:
                wr, hr = x / sin_a, x / cos_a
            else:
                wr, hr = x / cos_a, x / sin_a
        else:
            # fully constrained case: crop touches all 4 sides
            cos_2a = cos_a * cos_a - sin_a * sin_a
            wr = (w * cos_a - h * sin_a) / cos_2a
            hr = (h * cos_a - w * sin_a) / cos_2a

        rot_mat = cv2.getRotationMatrix2D((w / 2.0, h / 2.0), degree, 1.0)
        rot_mat[0, 2] += (wr - w) / 2.0
        rot_mat[1, 2] += (hr - h) / 2.0

        img = cv2.warpAffine(
            img, rot_mat, (int(round(wr)), int(round(hr))), flags=cv2.INTER_LINEAR
        )

        return img

    def __repr__(self):
        return f"Image rotation with angle range of (-{self.angle}, +{self.angle}) degrees."


class ColorJitter(object):
    def __init__(self, r=0.2):
        """Perturb color channels of a given image. Sample alpha in the range of (-r, r)
        and multiply (1 + alpha) to a color channel. The sampling is done independently
        for each channel.

        Parameters
        ----------
        r: float
            Range of color jitter ratio (-r ~ +r), max r = 1.0. Default: 0.2.
        """
        self.r = min(r, 1.0)

    def __call__(self, img):
        for i in range(img.shape[2]):
            alpha = random.uniform(-self.r, self.r)
            img[:, :, i] = img[:, :, i] * (1 + alpha)
        return img

    def __repr__(self):
        return "Random Color [Range {:.2f} - {:.2f}]".format(1 - self.r, 1 + self.r)


class GaussianBlur(object):
    def __init__(self, kernel_size=9, sigma_range=(0.1, 2.0), p=0.5):
        """
        Random Gaussian Blur of an image.

        Parameters
        ----------
        kernel_size: int
            Approximately 10% of the input image size (odd number). Default: 9.
        sigma_range: Tuple[float]
            Range of the sigmaX/Y to be uniformly sampled.
        p: float
            Probability of random augmentation. Default: 0.5.
        """
        self.kernel_size = (
            kernel_size
            if isinstance(kernel_size, tuple)
            else (kernel_size, kernel_size)
        )
        self.sigma_range = sigma_range
        self.p = p

    def __call__(self, img):
        if random.random() <= self.p:
            sigma = random.uniform(self.sigma_range[0], self.sigma_range[1])
            img = cv2.GaussianBlur(img, self.kernel_size, sigma)
        return img

    def __repr__(self):
        return "Random Gaussian Blur: [kernel={}] [sigma_range={}] [p={}]".format(
            self.kernel_size, self.sigma_range, self.p
        )


class ToTensor(object):
    def __init__(self, scale=10000):
        """
        Convert numpy array to tensor.

        Parameters
        ----------
        scale: int or float
            The max of pixel values of the given image.
            Default: 10000 for most remote sensing dataset.
        """
        self.scale = scale

    def __call__(self, img):
        img = img / self.scale
        img = np.clip(img, 0.0, 1.0)
        img = torch.from_numpy(img.transpose((2, 0, 1)))

        return img

    def __repr__(self):
        return "Numpy array image to tensor."


class Normalize_Std(object):
    def __init__(self, mean, std):
        """
        Normalize tensor image to N(0,1)

        Parameters
        ----------
        mean: Tuple[float]
            Mean pixel value of the image dataset.
        std: Tuple[float]
            Standard deviation of the pixel value of the image dataset.
        """
        self.mean = mean
        self.std = std
        self.normalize = transforms.Normalize(mean=mean, std=std)

    def __call__(self, img):
        img = self.normalize_post(img)
        return img

    def __repr__(self):
        return "Normalization to N(0,1), mean={} std={}".format(self.mean, self.std)
